package model;

import java.io.*;
import java.util.ArrayList;
import java.util.Random;

/**
 * Created by alex on 11/10/2016.
 */
public class BPNN {
    double[][] trainX;
    int[] trainY;

    double[][] validateX;
    int [] validateY;
    Layer layer0;
    Layer layer1;
    Layer layer2;
    int errocount = 0;

    public BPNN() {
        initTrainData();
        initValidaData();
        layer0 = new FullConnectLayer(784, 200);
        layer1 = new FullConnectLayer(200, 30);
        layer2 = new SoftmaxLayer(30, 8);
    }


    public void train() {
        int length = trainX.length;
        int count = 0;
        int maxEpoch = 100;
        for (int j = 0;j<maxEpoch;j++) {
            for (int i = 0; i < length; i++) {
                count++;
                forword(trainX[i], trainY[i]);
            }
            System.out.println("train number: " + count + "\ntest error rate is " + ((double) errocount) / length);
            errocount = 0;

            test();
//            System.out.println("layer0 norm: "+outputNorm(layer0.W));
//            System.out.println("layer1 norm: "+outputNorm(layer1.W));
//            System.out.println("layer2 norm: "+outputNorm(layer2.W));



        }
    }

    public void test(){
        int errorCount = 0;
        for (int i  = 0;i<validateX.length;i++){
            double[] output0 = layer0.output(validateX[i]);
            double[] output1 = layer1.output(output0);
            double[] output2 = layer2.output(output1);
            int yPred = findMax(output2);
            if (yPred!=validateY[i]){
                errorCount++;
            }
        }
        System.out.println("Validate Data Error rate: "+((double)errorCount)/validateX.length);
    }

    public void forword(double[] inputX, int y) {
        double[] output0 = layer0.output(inputX);
        double[] output1 = layer1.output(output0);
        double[] output2 = layer2.output(output1);
        backForword(inputX, output0, output1, output2, y);
    }

    public void backForword(double[] input, double[] output0, double[] output1, double[] output2, int y) {
        int yPred = findMax(output2);


        if (yPred != y) {
            errocount++;
        }

        //layer 2
        for (int i = 0; i < layer2.finOut; i++) {
            if (i == y) {
                layer2.delta[i] = output2[i] - 1.0;
            } else {
                layer2.delta[i] = output2[i] - 0.;
            }
        }

        for (int i = 0; i < layer2.finOut; i++) {
            for (int j = 0; j < layer2.finIn; j++) {
                layer2.deltaW[i][j] = layer2.delta[i] * output1[j];
            }
            layer2.deltab[i] = layer2.delta[i];
        }
        layer2.update(layer2.deltaW, layer2.deltab);

        //layer 1

        for (int i = 0; i < layer1.finOut; i++) {
            layer1.delta[i] = 0;
            for (int j = 0; j < layer2.finOut; j++) {
                layer1.delta[i] += layer2.W[j][i] * layer2.delta[j];
            }
            layer1.delta[i] *= (output1[i] * (1. - output1[i]));
        }

        for (int i = 0; i < layer1.finOut; i++) {
            for (int j = 0; j < layer1.finIn; j++) {
                layer1.deltaW[i][j] = layer1.delta[i] * output0[j];
            }
            layer1.deltab[i] = layer1.delta[i];
        }
        layer1.update(layer1.deltaW, layer1.deltab);


        //layer 0
        for (int i = 0; i < layer0.finOut; i++) {
            layer0.delta[i] = 0;
            for (int j = 0; j < layer1.finOut; j++) {
                layer0.delta[i] += layer1.W[j][i] * layer1.delta[j];
            }
            layer0.delta[i] *= (output0[i] * (1. - output0[i]));
        }

        for (int i = 0; i < layer0.finOut; i++) {
            for (int j = 0; j < layer0.finIn; j++) {
                layer0.deltaW[i][j] = layer0.delta[i] * input[j];
            }
            layer0.deltab[i] = layer0.delta[i];
        }
        layer0.update(layer0.deltaW, layer0.deltab);
    }

    private void initTrainData() {
        File file = new File("BPNN/src/dataset/train_data.txt");
        try {
            FileInputStream fileInputStream = new FileInputStream(file);
            BufferedReader br = new BufferedReader(new InputStreamReader(fileInputStream));
            ArrayList<double[]> X = new ArrayList<>();
            ArrayList<Integer> Y = new ArrayList<Integer>();
            for (String line = br.readLine(); line != null; line = br.readLine()) {
                line = line.trim();
                String[] pix = line.split(" ");
                if (pix.length != 785) {
                    System.out.println("read data length error");
                    return;
                }
                double[] xx = new double[784];
                for (int i = 0; i < 784; i++) {
                    xx[i] = (double) Integer.parseInt(pix[i]) / 255.;
                }
                int yy = Integer.parseInt(pix[784]);
                X.add(xx);
                Y.add(yy);
            }
            int size = X.size();
            this.trainX = (double[][]) X.toArray(new double[size][784]);
            this.trainY = new int[size];
            for (int i = 0; i < size; i++) {
                this.trainY[i] = Y.get(i);
            }
            shuffleArray(this.trainX,this.trainY);
            System.out.println("after shuffle");

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private void initValidaData(){
        File file = new File("BPNN/src/dataset/validation_data.txt");
        try {
            FileInputStream fileInputStream = new FileInputStream(file);
            BufferedReader br = new BufferedReader(new InputStreamReader(fileInputStream));
            ArrayList<double[]> X = new ArrayList<>();
            ArrayList<Integer> Y = new ArrayList<Integer>();
            for (String line = br.readLine(); line != null; line = br.readLine()) {
                line = line.trim();
                String[] pix = line.split(" ");
                if (pix.length != 785) {
                    System.out.println("read data length error");
                    return;
                }
                double[] xx = new double[784];
                for (int i = 0; i < 784; i++) {
                    xx[i] = (double) Integer.parseInt(pix[i]) / 255.;
                }
                int yy = Integer.parseInt(pix[784]);
                X.add(xx);
                Y.add(yy);
            }
            int size = X.size();
            this.validateX = (double[][]) X.toArray(new double[size][784]);
            this.validateY = new int[size];
            for (int i = 0; i < size; i++) {
                this.validateY[i] = Y.get(i);
            }
            shuffleArray(this.validateX,this.validateY);
            System.out.println("after shuffle");

        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    private int findMax(double[] array) {
        int index = 0;
        double max = array[0];
        for (int i = 1; i < array.length; i++) {
            if (array[i] > max) {
                index = i;
                max = array[i];
            }
        }
        return index;
    }

    private  void shuffleArray(double[][] x, int [] y)
    {
        // If running on Java 6 or older, use `new Random()` on RHS here
        Random rnd = new Random();
        for (int i = x.length - 1; i > 0; i--)
        {
            int index = rnd.nextInt(i + 1);
            // Simple swap
            double[] xx = x[index];
            x[index] = x[i];
            x[i] = xx;

            int yy = y[index];
            y[index] = y[i];
            y[i] = yy;
        }
    }


    public double outputNorm(double[][] matrix){
        int height = matrix.length;
        int width = matrix[0].length;
        double output = 0.;
        for (int i = 0;i<height;i++){
            for (int j = 0;j<width;j++){
                output+= matrix[i][j]*matrix[i][j];
            }
        }
        return Math.sqrt(output);
    }

}
